#!/usr/bin/env python

# Constants:
PACKAGE_NAME = 'annotations'
NODE_NAME = 'check_accuracy'
ANNOTATION_TOPIC = 'detections'
GROUND_TRUTH_TOPIC = 'ground_truth'

# -------
# Imports:
# -------
# ROS
import roslib; roslib.load_manifest(PACKAGE_NAME)
import rospy
from annotations.msg import AnnotationArray, Annotation
import cv

# Python
import Queue as queue

# -----------------
class CheckAccuracy:
# -----------------
	# ----------------
	def __init__(self):
	# ----------------
		# Initialize the node
		rospy.init_node(NODE_NAME)

		# Subscriptions
		rospy.Subscriber(GROUND_TRUTH_TOPIC, AnnotationArray, self.__groundTruthCallback)
		rospy.Subscriber(ANNOTATION_TOPIC, AnnotationArray, self.__annotationCallback)

		# Setup our data structures
		self.groundTruth = {} 				# Time --> ground truth message
		self.toProcess = queue.Queue()		# AnnotationArray messages that still need to be processed

	# ---------
	# Callbacks|
	# ----------------------------------
	def __groundTruthCallback(self, msg):
	# ----------------------------------
		# Add the ground truth message to the dictionary
		self.groundTruth[msg.header.stamp] = msg
	
	# ---------------------------------
	def __annotationCallback(self, msg):
	# ---------------------------------
		# Add the annotation message to the queue
		self.toProcess.put(msg)

	# -------
	# Private|
	# ---------------------------------------------------------------------------------
	def __processAnnotationArrayMessages(self, annotationArrayMsg, groundTruthArrayMsg):
	# ---------------------------------------------------------------------------------
		truePos = 0
		falsePos = 0
		falseNeg = 0		

		# Loop through the annotation array
		for annotation in annotationArrayMsg.annotations:
			# Compare the annotation to ALL of the ground truth rectangles
			foundMatch = False
			for groundTruth in groundTruthArrayMsg.annotations:
				foundMatch |= self.__compareRectanglesForDetection(annotation, groundTruth)
			
			if foundMatch:
				truePos += 1
			else:
				falsePos += 1

		falseNeg = len(groundTruthArrayMsg.annotations) - truePos 
		return (truePos, falsePos, falseNeg)	

	# ----------------------------------------------------------------
	def __compareRectanglesForDetection(self, annotation, groundTruth):
	# ----------------------------------------------------------------
		# Calculate the intersection of the rectangles
		intersection = self.__intersect(annotation, groundTruth)

		# If the rectangles don't intersect then it's not a valid detection
		if intersection == None:
			return False
		
		# Otherwise compare the area of the intersection -- if the are is >= 50% of the
		# area of the ground truth then count it as a detection
		x,y,w,h = intersection
		area = w * h
		groundTruthArea = groundTruth.width * groundTruth.height
		
		return (float(area) / float(groundTruthArea)) >= 0.5

	# ---------------------------
	def __intersect(self, r1, r2):
	# ---------------------------
		""" Calculate the intersection of two rectangles """
		tl = (max(r1.x, r2.x), max(r1.y, r2.y))
		br = (min(r1.x + r1.width, r2.x + r2.width), min(r1.y + r1.height, r2.y + r2.height))

		# Make sure the intersection is valid
		if tl[0] > br[0] or tl[1] > br[1]: 
			return None
		else:
			return (tl[0], tl[1], br[0] - tl[0], br[1] - tl[1])

	# ------
	# Public|
	# -----------
	def run(self):
	# -----------
		# Counters:
		totalTruePos = 0
		totalFalsePos = 0
		totalFalseNeg = 0
	
		# Main control loop
		# -----------------
		while not rospy.is_shutdown():
			try:
				# See if we have anything to process
				annotationMsg = self.toProcess.get(block = True, timeout = 0.1)

				try:
					# Find the corresponding ground truth message
					groundTruthMsg = self.groundTruth[annotationMsg.header.stamp]
	
					# Process the messages
					truePos, falsePos, falseNeg = self.__processAnnotationArrayMessages(annotationMsg, groundTruthMsg)
					print '%s: %i, %i, %i' % (str(annotationMsg.header.stamp), truePos, falsePos, falseNeg)

					# Update the counters
					totalTruePos += truePos
					totalFalsePos += falsePos
					totalFalseNeg += falseNeg

				except KeyError:
					# Put the message back in the queue 
					print 'No corresponding ground truth message: %s' % str(self.groundTruth)
					self.toProcess.put(annotationMsg)
					continue

			except queue.Empty:
				pass

		print '---'
		print 'Total true positives: %i' % totalTruePos
		print 'Total false positives: %i' % totalFalsePos
		print 'Total false negative: %i' % totalFalseNeg

# -----------------------
if __name__ == '__main__':
# -----------------------
	node = CheckAccuracy()
	node.run()
